import copy
import os

import cv2
import numpy as np
import torch

from htm.L46CorticalColumn import L46CorticalColumn
from htm.L46Inference import L46Inference
from Utils import Utils

# global Configurations
l4ColumnCount = 768
l4CellsPerColumn = 16
l4NumActiveMinicolumns = 3
l6ColumnCount = 110
l6CellsPerColumn = 110
l6NumOfLocModules = 10
l6BumpType = 'gaussian'


def getL6Params():
    locationModule2DParam = {
        'cellsPerAxis': l6ColumnCount,
        'scale': 40.0,
        'bumpOverlapMethod': 'probabilistic',
        'baselineCellsPerAxis': 6,
        'orientation': None,
        'connectionParams': {
            'columnCount': l6ColumnCount,
            'cellsPerColumn': l6CellsPerColumn,
            'initialPermanence': 1.0,
            'connectedPermanence': 0.5,
            'permanenceIncrement': 0.1,
            'permanenceDecrement': 0.0,
            'permanencePunishment': 0.0,
            'activationThreshold': 8,
            'minThreshold': 8,
            'sampleSize': 10,
            'maxSynapsesPerSegment': -1
        }
    }
    locationConfigs = []
    perModRange = float(90.0 if l6BumpType == "square" else 60.0 / float(l6NumOfLocModules))
    for i in range(l6NumOfLocModules):
        locModuleParamsCopy = copy.deepcopy(locationModule2DParam)  # copy the default
        orientationC = float(i) * perModRange + (perModRange / 2.0)
        locModuleParamsCopy['orientation'] = np.radians(orientationC)
        locationConfigs.append(locModuleParamsCopy)
    return locationConfigs


def getL4Params():
    params = {
        'columnCount': l4ColumnCount,
        'cellsPerColumn': l4CellsPerColumn,
        'reducedBasalThreshold': 8,
        'basalConnParams': {
            'columnCount': l4ColumnCount,
            'cellsPerColumn': l4CellsPerColumn,
            'initialPermanence': 1.0,
            'connectedPermanence': 0.5,
            'permanenceIncrement': 0.1,
            'permanenceDecrement': 0.0,
            'permanencePunishment': 0.0,
            'activationThreshold': 8,
            'minThreshold': 10,
            'sampleSize': 10,
            'maxSynapsesPerSegment': -1
        },
        'apicalConnParams': {
            'columnCount': l4ColumnCount,
            'cellsPerColumn': l4CellsPerColumn,
            'initialPermanence': 1.0,
            'connectedPermanence': 0.5,
            'permanenceIncrement': 0.1,
            'permanenceDecrement': 0.0,
            'permanencePunishment': 0.0,
            'activationThreshold': 8,
            'minThreshold': 10,
            'sampleSize': 10,
            'maxSynapsesPerSegment': -1
        }
    }
    return params


def getL46CortColParams():
    params = {
        'bumpType': l6BumpType,
        'numOfLocModules': 10,
        'thresholds': -1
    }
    return params


def getL46InferenceParams():
    params = {
        'noiseFactor': 0,
        'moduleNoiseFactor': 0,
        'maxSettlingTime': 10,
        'maxTraversals': 4,
        'nextMonitorToken': 1,
        'randomLocation': False,
        'useNoise': False,
        'noisyTrainingTime': 1,
        'waitForSettle': False,
        'numSensations': 9
    }
    return params


def runExperiment():
    if not os.path.exists("trained_models/"):
        os.makedirs("trained_models/")
    np.random.seed(1)  # The only place to set seed
    imgPath = "../visionBodyLearn/outputImages/resaved_temp0.png"
    origScene1 = cv2.imread(imgPath)
    edgeScene1 = Utils.getEdgeImg(origScene1)
    scene1 = Utils.encodeScenarioUsingImage(origScene1, edgeScene1)
    l46CortColParams = getL46CortColParams()
    l6Params = getL6Params()
    l4Params = getL4Params()
    l46CorticalColumn = L46CorticalColumn(l46CortColParams, l6Params, l4Params)
    l46InferenceParams = getL46InferenceParams()
    l46Inference = L46Inference(l46CorticalColumn, l46InferenceParams)
    for objectDescription in [scene1]:
        print('learning Object:', objectDescription['name'])
        l46Inference.learnObject(objectDescription)
    torch.save(l46Inference, 'trained_models/piunExp.m')
    torch.save(edgeScene1, 'trained_models/edgeScene1.m')


if __name__ == "__main__":
    runExperiment()
